import sounddevice as sd
import soundfile as sf
import webrtcvad
import numpy as np
import whisper
import threading
import time

class SimpleWakeupDetector:
    def __init__(self, wake_words=["你好", "小助手", "机器人"], sensitivity=0.8):
        """
        简单的唤醒词检测器，基于VAD和Whisper识别
        :param wake_words: 唤醒词列表
        :param sensitivity: 灵敏度，0-1之间
        """
        self.wake_words = wake_words
        self.sensitivity = sensitivity
        self.vad = webrtcvad.Vad(2)  # VAD模式2
        self.whisper_model = None
        self.is_listening = False
        
    def load_whisper_model(self):
        """延迟加载Whisper模型"""
        if self.whisper_model is None:
            print("加载Whisper模型...")
            self.whisper_model = whisper.load_model("tiny")
            
    def record_audio(self, duration=2, samplerate=16000):
        """录制短音频"""
        print("录制音频中...")
        audio = sd.rec(int(duration * samplerate), samplerate=samplerate, channels=1, dtype='int16')
        sd.wait()
        return audio
        
    def detect_wake_word(self, audio, samplerate=16000):
        """检测唤醒词"""
        try:
            # 保存临时音频文件
            temp_file = "temp_wakeup.wav"
            sf.write(temp_file, audio, samplerate)
            
            # 使用Whisper识别
            if self.whisper_model is None:
                self.load_whisper_model()
                
            result = self.whisper_model.transcribe(temp_file, language="zh")
            text = result["text"].strip().lower()
            
            print(f"识别结果: {text}")
            
            # 检查是否包含唤醒词
            for wake_word in self.wake_words:
                if wake_word.lower() in text:
                    return True, wake_word
                    
            return False, None
            
        except Exception as e:
            print(f"唤醒词检测异常: {e}")
            return False, None
            
    def wait_for_wakeup(self, callback=None):
        """
        等待唤醒词
        :param callback: 检测到唤醒词后执行的回调函数
        """
        print(f"等待唤醒词: {self.wake_words}")
        print("请说出唤醒词...")
        
        while True:
            try:
                # 录制2秒音频
                audio = self.record_audio(duration=2)
                
                # 检测唤醒词
                detected, wake_word = self.detect_wake_word(audio)
                
                if detected:
                    print(f"检测到唤醒词: {wake_word}")
                    if callback:
                        callback()
                    break
                    
                time.sleep(0.5)  # 短暂休息
                
            except KeyboardInterrupt:
                print("用户中断")
                break
            except Exception as e:
                print(f"唤醒检测异常: {e}")
                time.sleep(1)

def wait_for_wakeup(callback=None, wake_words=["你好", "小助手", "机器人"]):
    """
    简化的唤醒词检测接口
    """
    detector = SimpleWakeupDetector(wake_words=wake_words)
    detector.wait_for_wakeup(callback) 